﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.ComponentModel;
using System.IO;
using System.Runtime.InteropServices;
using System.Drawing;

namespace ThumbcacheHelper
{
    class CacheCruncher
    {
        byte[] MAGIC_ENTRY = new byte[] { 0x43, 0x4D, 0x4D, 0x4D }; // CMMM

        byte[] MAGIC_BMP = new byte[] { 0x42, 0x4D }; // BM
        byte[] MAGIC_JPG = new byte[] { 0xFF, 0xD8, 0xFF, 0xE0 };
        byte[] MAGIC_PNG = new byte[] { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };

        [StructLayout(LayoutKind.Explicit)]
        struct FileHeader
        {
            [FieldOffset(0)]
            public int magic;

            [FieldOffset(4)]
            public uint version;

            [FieldOffset(8)]
            public uint type;

            [FieldOffset(12)]
            public uint firstEntryOffset;

            [FieldOffset(16)]
            public uint nextEntryOffset;

            [FieldOffset(20)]
            public uint entryCount;
        }

        [StructLayout(LayoutKind.Explicit)]
        struct Entry_7
        {
            [FieldOffset(0)]
            public int magic;

            [FieldOffset(4)]
            public uint size;

            [FieldOffset(8)]
            public long hash;

            [FieldOffset(16)]
            public uint filenameLength;

            [FieldOffset(20)]
            public uint paddingSize;

            [FieldOffset(24)]
            public uint dataSize;

            [FieldOffset(28)]
            public uint unknown;

            [FieldOffset(32)]
            public long dataChecksum;

            [FieldOffset(40)]
            public long headerChecksum;
        }

        [StructLayout(LayoutKind.Explicit)]
        struct Entry_Vista
        {
            [FieldOffset(0)]
            public int magic;

            [FieldOffset(4)]
            public uint size;

            [FieldOffset(8)]
            public long hash;

            [FieldOffset(16)]
            public long extension;

            [FieldOffset(24)]
            public uint filenameLength;

            [FieldOffset(28)]
            public uint paddingSize;

            [FieldOffset(32)]
            public uint dataSize;

            [FieldOffset(36)]
            public uint unknown;

            [FieldOffset(40)]
            public long dataChecksum;

            [FieldOffset(48)]
            public long headerChecksum;
        }

        [StructLayout(LayoutKind.Explicit)]
        struct IdxHeader
        {
            [FieldOffset(0)]
            public uint magic;

            [FieldOffset(4)]
            public uint version;

            [FieldOffset(8)]
            public uint unknown;

            [FieldOffset(12)]
            public uint usedEntries;

            [FieldOffset(16)]
            public uint entries;
        }

        [StructLayout(LayoutKind.Explicit)]
        struct IdxEntry_Vista
        {
            [FieldOffset(0)]
            public long hash;

            [FieldOffset(8)]
            public long modified;

            [FieldOffset(16)]
            public uint flags;

            [FieldOffset(20)]
            public uint offset32;

            [FieldOffset(24)]
            public uint offset96;

            [FieldOffset(28)]
            public uint offset256;

            [FieldOffset(32)]
            public uint offset1024;

            [FieldOffset(36)]
            public uint offsetsr;
        }

        [StructLayout(LayoutKind.Explicit)]
        struct IdxEntry_7
        {
            [FieldOffset(0)]
            public long hash;

            [FieldOffset(8)]
            public uint flags;

            [FieldOffset(12)]
            public uint offset32;

            [FieldOffset(16)]
            public uint offset96;

            [FieldOffset(20)]
            public uint offset256;

            [FieldOffset(24)]
            public uint offset1024;

            [FieldOffset(28)]
            public uint offsetsr;
        }

        readonly int fileHeaderSize = Marshal.SizeOf(typeof(FileHeader));
        readonly int entrySize_7 = Marshal.SizeOf(typeof(Entry_7));
        readonly int entrySize_Vista = Marshal.SizeOf(typeof(Entry_Vista));
        readonly int idxHeaderSize = Marshal.SizeOf(typeof(IdxHeader));
        readonly int idxEntrySize_7 = Marshal.SizeOf(typeof(IdxEntry_7));
        readonly int idxEntrySize_Vista = Marshal.SizeOf(typeof(IdxEntry_Vista));

        BackgroundWorker _bgw = new BackgroundWorker();
        Dictionary<FileId, String> _if;
        FrmMain _parent;
        bool useSeeker = false;

        public CacheCruncher(FrmMain parent, Dictionary<FileId, String> inputFiles)
        {
            _parent = parent;
            _if = inputFiles;

            _bgw.ProgressChanged += new ProgressChangedEventHandler(_bgw_ProgressChanged);
            _bgw.DoWork += new DoWorkEventHandler(_bgw_DoWork);
            _bgw.RunWorkerCompleted += new RunWorkerCompletedEventHandler(_bgw_RunWorkerCompleted);

            _bgw.WorkerSupportsCancellation = true;
            _bgw.WorkerReportsProgress = true;
        }

        /// <summary>
        /// Cancel the CacheCruncher process
        /// </summary>
        public void Cancel()
        {
            _bgw.CancelAsync();
        }

        void _bgw_RunWorkerCompleted(object sender, RunWorkerCompletedEventArgs e)
        {
            _parent.UpdateDataGridView();
            _parent.log("Done.");
        }

        void _bgw_DoWork(object sender, DoWorkEventArgs e)
        {
            Dictionary<FileId, String> inputFiles = (Dictionary<FileId, String>)e.Argument;
            EdbSeeker edbSeeker = null;

            if (inputFiles.ContainsKey(FileId.WinEdb))
            {
                _bgw.ReportProgress(0, new FeedbackMessage("Preparing Windows Search Index for lookup..."));
                edbSeeker = new EdbSeeker();
                string prepResult = edbSeeker.Prepare(inputFiles[FileId.WinEdb]);
                if (prepResult == null)
                {
                    edbSeeker.CreateIndex_Vista();

                    _bgw.ReportProgress(0, new FeedbackMessage("Done", true));
                    useSeeker = true;
                }
                else
                {
                    _bgw.ReportProgress(0, new FeedbackMessage(prepResult, true, Color.IndianRed));
                }
            }


            // Take each fileid in the order of the Enum - do the thumbcaches before the idx
            foreach (var fileId in Enum.GetValues(typeof(FileId)))
            {
                if (_bgw.CancellationPending)
                    break;

                FileId thisFileId = (FileId)fileId;
                if (inputFiles.ContainsKey(thisFileId)) // Supplied that file?
                {
                    if (thisFileId < FileId.ThumbsIdx) // Is a thumbcache
                    {
                        if (!string.IsNullOrWhiteSpace(inputFiles[thisFileId]))
                            processThumbCache(inputFiles[thisFileId], edbSeeker, thisFileId);
                    }
                    else if (thisFileId == FileId.ThumbsIdx)
                        processIdx(inputFiles[thisFileId]);
                }
            }
        }

        void _bgw_ProgressChanged(object sender, ProgressChangedEventArgs e)
        {
            switch (e.ProgressPercentage)
            {
                case 0: // Log message
                    FeedbackMessage fm = (FeedbackMessage)e.UserState;
                    _parent.log(fm.Message, fm.Colour, fm.AddBlankLink);
                    break;
                case 1: // Add Entry
                    _parent.AddEntry((EntryData)e.UserState);
                    break;
                case 2: // Set max for progress bar
                    _parent.SetProgressMax((int)e.UserState);
                    break;
                case 3: // Set status bar msg
                    _parent.SetStatus((string)e.UserState);
                    break;
                case 4: // Reset progress bar
                    _parent.ResetProgressBar();
                    break;
                case 5: // Increment progress bar only
                    _parent.IncrementProgressBar((int)e.UserState);
                    break;
            }
        }

        void processIdx(string idxFile)
        {
            _bgw.ReportProgress(0, new FeedbackMessage(idxFile));
            _bgw.ReportProgress(3, idxFile);
            _bgw.ReportProgress(4);

            using (BinaryReader br = new BinaryReader(File.OpenRead(idxFile)))
            {
                int unused = 0;
                int used = 0;

                byte[] b = br.ReadBytes(idxHeaderSize);
                GCHandle gch = GCHandle.Alloc(b, GCHandleType.Pinned);
                IdxHeader idxHeader = (IdxHeader)Marshal.PtrToStructure(gch.AddrOfPinnedObject(), typeof(IdxHeader));
                gch.Free();

                _bgw.ReportProgress(2, checked((int)idxHeader.entries));
                _bgw.ReportProgress(0, new FeedbackMessage(getOSVersion(idxHeader.version)));
                _bgw.ReportProgress(0, new FeedbackMessage(String.Format("Expecting {0} entries, {1} used.", idxHeader.entries, idxHeader.usedEntries)));

                br.BaseStream.Seek(24, SeekOrigin.Begin);

                for (int i = 0; i < idxHeader.entries; i++)
                {
                    if (_bgw.CancellationPending)
                        return;

                    long entryOffset = br.BaseStream.Position;
                    IdxEntry_Vista e = new IdxEntry_Vista();

                    if (idxHeader.version == 0x14) // Vista
                    {
                        b = br.ReadBytes(idxEntrySize_7);
                        gch = GCHandle.Alloc(b, GCHandleType.Pinned);
                        e = (IdxEntry_Vista)Marshal.PtrToStructure(gch.AddrOfPinnedObject(), typeof(IdxEntry_Vista));
                        gch.Free();
                    }
                    else if (idxHeader.version == 0x15) // 7
                    {
                        b = br.ReadBytes(idxEntrySize_7);
                        gch = GCHandle.Alloc(b, GCHandleType.Pinned);
                        IdxEntry_7 idxEntry = (IdxEntry_7)Marshal.PtrToStructure(gch.AddrOfPinnedObject(), typeof(IdxEntry_7));
                        gch.Free();

                        e.flags = idxEntry.flags;
                        e.hash = idxEntry.hash;
                        e.modified = -1;
                        e.offset1024 = idxEntry.offset1024;
                        e.offset256 = idxEntry.offset256;
                        e.offset32 = idxEntry.offset32;
                        e.offset96 = idxEntry.offset96;
                        e.offsetsr = idxEntry.offsetsr;
                    }
                    else
                        return;

                    byte[] hashAsBytes = BitConverter.GetBytes(e.hash);
                    Array.Reverse(hashAsBytes);
                    string hashAsString = Utils.ByteArrayToByteString(hashAsBytes).ToLower();

                    _parent.SetIdxOffset(FileId.Thumbs32, hashAsString, e.offset32, entryOffset);
                    _parent.SetIdxOffset(FileId.Thumbs96, hashAsString, e.offset96, entryOffset);
                    _parent.SetIdxOffset(FileId.Thumbs256, hashAsString, e.offset256, entryOffset);
                    _parent.SetIdxOffset(FileId.Thumbs1024, hashAsString, e.offset1024, entryOffset);

                    _bgw.ReportProgress(5, 1);
                }

                _bgw.ReportProgress(0, new FeedbackMessage(String.Format("Processed {0} entries: {1} used, {2} unused.", used + unused, used, unused), true));
            }
        }

        void processThumbCache(string thumbcacheFile, EdbSeeker edbSeeker, FileId FileId)
        {
            _bgw.ReportProgress(0, new FeedbackMessage(thumbcacheFile));
            _bgw.ReportProgress(3, thumbcacheFile);
            _bgw.ReportProgress(4);

            using (BinaryReader br = new BinaryReader(File.OpenRead(thumbcacheFile)))
            {
                int processed = 0;

                byte[] b = br.ReadBytes(fileHeaderSize);
                GCHandle gch = GCHandle.Alloc(b, GCHandleType.Pinned);
                FileHeader fileHeader = (FileHeader)Marshal.PtrToStructure(gch.AddrOfPinnedObject(), typeof(FileHeader));
                gch.Free();

                _bgw.ReportProgress(2, checked((int)fileHeader.entryCount));
                _bgw.ReportProgress(0, new FeedbackMessage(String.Format("{0}, {1}", getOSVersion(fileHeader.version), getCacheType(fileHeader.version, fileHeader.type))));
                _bgw.ReportProgress(0, new FeedbackMessage(String.Format("Expecting {0} thumbnails.", fileHeader.entryCount)));

                br.BaseStream.Seek(fileHeader.firstEntryOffset, SeekOrigin.Begin);

                for (int i = 0; i < fileHeader.entryCount; i++)
                {
                    long startOfEntry = br.BaseStream.Position; // Need this later

                    if (_bgw.CancellationPending)
                        return;

                    EntryData entryData = new EntryData();
                    entryData.FileId = FileId;
                    entryData.FilePath = thumbcacheFile;
                    entryData.Offset = br.BaseStream.Position;

                    try
                    {
                        if (i == 5808)
                            Console.WriteLine();

                        Entry_Vista e = new Entry_Vista();

                        if (fileHeader.version == 0x14) // Vista
                        {
                            b = Utils.AssuredReadBytes(br, entrySize_Vista);
                            gch = GCHandle.Alloc(b, GCHandleType.Pinned);
                            Entry_Vista entry = (Entry_Vista)Marshal.PtrToStructure(gch.AddrOfPinnedObject(), typeof(Entry_Vista));
                            gch.Free();

                            e = entry;
                        }
                        else if (fileHeader.version == 0x15) // 7
                        {
                            b = Utils.AssuredReadBytes(br, entrySize_7);
                            gch = GCHandle.Alloc(b, GCHandleType.Pinned);
                            Entry_7 entry = (Entry_7)Marshal.PtrToStructure(gch.AddrOfPinnedObject(), typeof(Entry_7));
                            gch.Free();

                            e.dataChecksum = entry.dataChecksum;
                            e.dataSize = entry.dataSize;
                            e.extension = -1L;
                            e.filenameLength = entry.filenameLength;
                            e.hash = entry.hash;
                            e.headerChecksum = entry.headerChecksum;
                            e.magic = entry.magic;
                            e.paddingSize = entry.paddingSize;
                            e.size = entry.size;
                            e.unknown = entry.unknown;
                        }
                        else
                        {
                            return;
                        }

                        // Check for last entry
                        if (startOfEntry + e.size >= br.BaseStream.Length)
                        {
                            break;
                        }

                        entryData.EntryHash = e.hash.ToString("X16").ToLower();
                        entryData.DataSize = (long)e.dataSize;
                        entryData.DataChecksum = e.dataChecksum.ToString("X16").ToLower();
                        entryData.HeaderChecksum = e.headerChecksum.ToString("X16").ToLower();

                        if (e.filenameLength > 0)
                        {
                            b = Utils.AssuredReadBytes(br, checked((int)e.filenameLength));
                            entryData.FileName = Encoding.Unicode.GetString(b);


                            if (useSeeker)
                            {
                                //if (fileHeader.version == 0x14) // Vista
                                //{
                                entryData.SourceFileName = edbSeeker.LookupHash_Vista(entryData.FileName, e.hash);
                                //}
                                //else if (fileHeader.version == 0x15) // 7
                                //{
                                //    entryData.SourceFileName = edbSeeker.LookupHash_Vista(entryData.FileName);
                                //}
                            }


                            if (e.dataChecksum != 0)
                            {
                                br.BaseStream.Seek(e.paddingSize, SeekOrigin.Current);

                                if (e.dataSize > 0)
                                {
                                    b = Utils.AssuredReadBytes(br, checked((int)e.dataSize));

                                    if (Utils.ArrayStartsWith(b, MAGIC_BMP))
                                    {
                                        entryData.Type = ImageType.BMP;
                                    }
                                    else if (Utils.ArrayStartsWith(b, MAGIC_JPG))
                                    {
                                        entryData.Type = ImageType.JPG;
                                    }
                                    else if (Utils.ArrayStartsWith(b, MAGIC_PNG))
                                    {
                                        entryData.Type = ImageType.PNG;
                                    }

                                    entryData.setThumbBytes(b);
                                }
                            }
                        }
                        else
                        {
                            if (fileHeader.version == 0x14) // Vista
                                br.BaseStream.Seek(e.size - entrySize_Vista, SeekOrigin.Current);
                            else if (fileHeader.version == 0x15)
                                br.BaseStream.Seek(e.size - entrySize_7, SeekOrigin.Current);
                        }

                        _bgw.ReportProgress(1, entryData);
                        processed++;

                        br.BaseStream.Seek(startOfEntry + e.size, SeekOrigin.Begin);                        
                    }
                    catch (Exception ex)
                    {
                        _bgw.ReportProgress(0, new FeedbackMessage(String.Format("Error processing record @offset {0}: {1}", entryData.Offset, ex.Message), Colour: Color.IndianRed));
                    }
                }

                _bgw.ReportProgress(0, new FeedbackMessage(String.Format("Processed {0} records.", processed), true));
            }
        }        

        /// <summary>
        /// Starts the analysis of the files on a background thread.
        /// </summary>
        public void Start()
        {
            _bgw.RunWorkerAsync(_if);
        }

        public bool IsRunning()
        {
            return _bgw.IsBusy;
        }

        private CacheType getCacheType(uint version, uint type)
        {
            if (version == 0x14 || version == 0x15)
                return (CacheType)type;
            else
                return CacheType.Unknown;
        }

        private string getOSVersion(uint version)
        {
            string os = "Unknown";

            switch (version)
            {
                case 0x14:
                    os = "Windows Vista";
                    break;
                case 0x15:
                    os = "Windows 7";
                    break;
                case 0x1A:
                    os = "Windows 8";
                    break;
                case 0x1C:
                    os = "Windows 8 Version 2";
                    break;
            }

            return os;
        }
    }
}
